#include <sstream>
#include <iomanip>
#include <opencv2/dnn.hpp>

#include "base_util/utils.h"
#include "postprocess_op.h"
#include "preprocess_op.h"
#include "ocr_rec.h"


namespace ai {


typedef unsigned short half;
typedef unsigned int uint;

uint as_uint(const float x) {
    return *(uint*)&x;
}

float as_float(const uint x) {
    return *(float*)&x;
}

float half_to_float(half x) { // IEEE-754 16-bit floating-point format (without infinity): 1-5-10, exp-15, +-131008.0, +-6.1035156E-5, +-5.9604645E-8, 3.311 digits
    // printf("1\n");
    const uint e = (x&0x7C00)>>10; // exponent
    // printf("2\n");
    const uint m = (x&0x03FF)<<13; // mantissa
    // printf("3\n");
    const uint v = as_uint((float)m)>>23; // evil log2 bit hack to count leading zeros in denormalized format
    // printf("4\n");
    return as_float((x&0x8000)<<16 | (e!=0)*((e+112)<<23|m) | ((e==0)&(m!=0))*((v-37)<<23|((m<<(150-v))&0x007FE000))); // sign : normalized : denormalized
}

static void dump_tensor_attr(rknn_tensor_attr *attr) {
  printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
          "zp=%d, scale=%f\n",
          attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
          attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
          get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}

static float DeQnt2F32(int8_t qnt, int zp, float scale) {
  return ((float)qnt - (float)zp) * scale;
}

RKOcrRecPredictor::RKOcrRecPredictor(RKPackPredictor* model, RKModelManager* manager, LogInfo *lg):
  RKPredictor(model, manager, lg) { }


int RKOcrRecPredictor::start_postprocess_thread_imp() {
  int ret = 0;
  std::pair<std::pair<int,int>,int> cur_block;
  // stop_flag2 = common::CameraGrabing;
  // while (stop_flag2 == common::CameraGrabing) {

  //   run_mtx2.lock();
  //   if (post_data2.empty()) {run_mtx2.unlock(); std::this_thread::sleep_for(std::chrono::milliseconds(3));continue;}
  //   cur_block = post_data2.front();
  //   post_data2.pop();
  //   run_mtx2.unlock();

  //   int loop_idx = cur_block.first.second;
  //   std::vector<BaseInfo*> det_infos;
  //   ret = det_hb_->PostProcess(det_infos, cur_block.first.first, this, cur_block.second);  //处理图片
  //   run_mtx3.lock();
  //   out_data3[loop_idx].push(det_infos);  //将处理后的数据返回到det_infos,加入队列中
  //   run_mtx3.unlock();
  // }
  // stop_flag2 = common::CameraOpened;
  return 0;
}

int RKOcrRecPredictor::RunDet(stream::ImageBlob* blob, std::vector<BaseInfo*>& det_infos) {
  int ret;

  stream::ImageBlob img_blob(stream::ImageBlobMode_BGR);
  if (!mdl_rk->cfg->transforms->run(blob->img, img_blob, mdl_rk->cfg)) { 
    printf("transforms->run fail \n");
    return model_image_channels_check_error;
  }
  // std::cout << "img_blob.img w:" << img_blob.img.cols << " h:" << img_blob.img.rows << std::endl;
  // cv::imwrite("padimg.jpg", img_blob.img);


  rknn_tensor_attr output_attrs[mdl_rk->output_count];
  memset(output_attrs, 0, sizeof(output_attrs));
  for (int i = 0; i < mdl_rk->output_count; i++) {
    output_attrs[i].index = i;
    ret = rknn_query(mdl_rk->ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
    // dump_tensor_attr(&(output_attrs[i]));
  } 

  mdl_rk->inputs[0].buf = (void *)(img_blob.img.data);
  rknn_inputs_set(mdl_rk->ctx, mdl_rk->input_count, &mdl_rk->inputs[0]);

  rknn_core_mask core_mask = RKNN_NPU_CORE_0;
  if(mdl_rk->cfg->gpu_id==1) { core_mask = RKNN_NPU_CORE_1; }
  if(mdl_rk->cfg->gpu_id==2) { core_mask = RKNN_NPU_CORE_2; }
  ret = rknn_set_core_mask(mdl_rk->ctx, core_mask);
  ret = rknn_run(mdl_rk->ctx, NULL);
  if (ret != 0) {printf("rknn_run fail %d", ret); return -1;}
  ret = rknn_outputs_get(mdl_rk->ctx, mdl_rk->output_count, &mdl_rk->outputs[0], NULL);
  if (ret != 0) {printf("rknn_outputs_get fail %d", ret); return -1;}
 
  // 后处理部分
  std::vector<float> out_scales;
  std::vector<int32_t> out_zps;
  for (int i = 0; i < mdl_rk->output_count; ++i)
  {
    out_scales.push_back(output_attrs[i].scale);
    out_zps.push_back(output_attrs[i].zp);
  }

  std::vector<std::vector<float>> output_tensors(mdl_rk->output_count, std::vector<float>(I_SIZE_CARPLATE_REC));
  for (int j = 0; j < mdl_rk->output_count; j++)
  {

    int8_t *pblob = (int8_t *)mdl_rk->outputs[j].buf;
    for(int i=0;i<output_attrs[j].n_elems;i++)
    {
        output_tensors[j][i] = DeQnt2F32(pblob[i], out_zps[j], out_scales[j]);             

    }
    // free(pblob);
    // pblob=NULL;

  }
  // std::cout << "mdl_rk->output_count:" << mdl_rk->output_count << std::endl;
  // std::cout << output_tensors[0][0] << " " << output_tensors[0][1] << " " << output_tensors[0][2] << " " << output_tensors[0][3] << " " << output_tensors[0][4] << " " << output_tensors[0][5] << " " << output_tensors[0][6] << " " << output_tensors[0][7] << " " << output_tensors[0][8] << " " << output_tensors[0][9] << " " << output_tensors[0][10] << std::endl; 
  // std::cout << output_tensors[1][0] << " " << output_tensors[1][1] << " " << output_tensors[1][2] << " " << output_tensors[1][3] << " " << output_tensors[1][4] << " " << output_tensors[1][5] << " " << output_tensors[1][6] << " " << output_tensors[1][7] << " " << output_tensors[1][8] << " " << output_tensors[1][9] << " " << output_tensors[1][10] << std::endl; 

  GetConvDetectionResult(mdl_rk->cfg, output_tensors, out_zps, out_scales, det_infos);

  rknn_outputs_release(mdl_rk->ctx, mdl_rk->output_count, &mdl_rk->outputs[0]);

  return 0;

}
static inline float fast_exp(float x)
{
    // return exp(x);
    union
    {
        uint32_t i;
        float f;
    } v;
    v.i = (12102203.1616540672 * x + 1064807160.56887296);
    return v.f;
}

static std::vector<std::string> color_list = {"黑色","蓝色","绿色","白色","黄色"};
int RKOcrRecPredictor::GetConvDetectionResult(ai::ModelConfig *cfg, std::vector<std::vector<float>>& output_tensors, std::vector<int> &qnt_zp, std::vector<float> &qnt_scale, std::vector<BaseInfo*>& det_infos) {

    float max = output_tensors[1][0];
    int color_index = 0;
    int output_size_color = 5; //颜色5种
    for (int i = 1 ; i<output_size_color;i++)
    {
        if(output_tensors[1][i]>max)
        {
            max = output_tensors[1][i];
            color_index = i;
        }
    }
    // printf("[liufeng debug] +\n");
    std::string plate_color = color_list[color_index];
    // printf("plate_color:%s\n",plate_color.c_str());
    int time_step = 21;   //时间步长
    int num_char=78; //字符类别数
    int output_size = time_step*num_char;
    std::string plate_no;
    float plate_no_conf = 1;
    
    int pre_index = 0;
    float* prob = &output_tensors[0][0];
    for(int i = 0; i<time_step;i++)
    {
        float *row_ptr = prob+num_char*i;
        int max_value = row_ptr[0];
        int max_index = 0;
        float fexp = 0;
        float fexps[num_char-1]={0};
        for(int j = 1; j<num_char;j++)
        {   
            // printf("%.4f\n",row_ptr[j]);
            fexp+=fast_exp(row_ptr[j]);
            // if(i == 0 && j < 4){printf("fast_exp(row_ptr[j]):%.4f\n",fast_exp(row_ptr[j]));}
            fexps[j-1] = fast_exp(row_ptr[j]);
            if (row_ptr[j]>max_value)
            {
                max_value= row_ptr[j];
                max_index = j;
            }
        }

        for(int j = 0; j<num_char-1;j++)
        {
            fexps[j]/=fexp;
            // printf("%.4f\n",fexps[j]);
        }
        
        if (max_index!=0 && max_index!=pre_index)
        {
            plate_no+=cfg->label_list[max_index];
            //printf("%.4f\n",fexp);
            //printf("plate_no conf:%.4f\n",fexps[max_index-1]);
            plate_no_conf*=fexps[max_index-1];
        }

        pre_index= max_index;
    }

    // printf("plate_no_conf:%f\n",plate_no_conf);
    // printf("plate_no:%s\n",plate_no.c_str());

    TextRecInfo* di = new TextRecInfo(plate_no_conf, plate_no, plate_color);
    det_infos.emplace_back(di);

  return 0;
}


}  // namespace hb
